#include <track.h>
#include <defs.h>
#include <user.h>
#include <string.h>
#include <printer.h>
#include <io.h>
#include <model.h>
#include <iobuf.h>
#include <utils.h>

/**
 * Uses BFS to find all sensor-neighbours of a given node. This will be
 * used to form multiple predictions.
 */
void BFS(int a, track_model *graph, track_edge *nbrs) {
  int nbrCount = 0;
  int dist[MAX_NODES];
  int seen[MAX_NODES];
  iobuf Q;
  IobufInit(&Q);

  int i, u;
  for (i = 0; i < MAX_NODES; i++) {
    dist[i] = INFINITY;
    seen[i] = FALSE;
  }
  dist[a] = 0;

  IobufEnqueue(&Q, a);
  while (!IobufEmpty(&Q)) {
    IobufDequeue(&Q, &u);
    seen[u] = TRUE;

    int edges;
    track_node current = graph->nodes[u];

    if (current.type == NODE_SWITCH)
      edges = 3;
    else if (current.type == NODE_SENSOR) {
      edges = 2;
      nbrs[nbrCount].dest = u;
      nbrs[nbrCount].distance = dist[u];
      nbrCount++;
      if (u != a) {
        continue;
      }
    } else if (current.type == NODE_STOP)
      edges = 1;
   
    for (i = 0; i < edges; i++) {
      int v = current.edges[i].dest;
      if (seen[v]) continue;
      int alt = dist[u] + current.edges[i].distance;
      if (alt < dist[v])
        dist[v] = alt;
      IobufEnqueue(&Q, v);
    }
  }
  nbrs[nbrCount].dest = UNDEFINED;
}

int MyNextLandmark(track_model *tm, int prev, int curr, track_edge *next) {
  switch (tm->nodes[curr].type) {
  case NODE_SWITCH:
    if (prev == tm->nodes[curr].sw.behind.dest) {
      // get switch position
      int switch_number = tm->nodes[curr].id2;
      int switch_direction = tm->nodes[curr].sw.set;
      switch (switch_direction) {
      case 'C':
        *next = tm->nodes[curr].sw.ahead[0];
        break;
      case 'S':
        *next = tm->nodes[curr].sw.ahead[1];
        break;
      }
    } else {
      *next = tm->nodes[curr].sw.behind;
    }
    break;
  case NODE_SENSOR:
    if (prev == tm->nodes[curr].se.ahead.dest)
      *next = tm->nodes[curr].se.behind;
    else if (prev == tm->nodes[curr].se.behind.dest)
      *next = tm->nodes[curr].se.ahead;
    else
      return FALSE;
    break;
  case NODE_STOP:
    *next = tm->nodes[curr].st.ahead;
    break;
  default:
    Panic("MyNextLandmark", "invalid track node type");
    break;
  }
  return TRUE;
}

void MySensorNextLandmark(
    track_model *tm, int sensor_id, int sensor_direction, track_edge *next) {
  track_node curr = tm->nodes[sensor_id];
  if (curr.type != NODE_SENSOR)
    Panic("MySensorNextLandmark", "invalid node - must be sensor");
  switch (sensor_direction) {
  case SENSOR_FORWARD:
    *next = curr.se.ahead;
    break;
  case SENSOR_BACK:
    *next = curr.se.behind;
    break;
  default:
    Panic("MySensorNextLandmark", "invalid sensor direction");
  }
}

void MyNextSensors(track_model *tm, int sensor_id, int sensor_direction,
    int n, track_edge *next_sensors) {
  int i;
  int prev = sensor_id;
  track_edge next_sensor, next;
  MySensorNextLandmark(tm, sensor_id, sensor_direction, &next_sensor);
  int curr = next_sensor.dest;
  // get next landmark in direction
  for (i = 0; i < n; i++) {
    while (tm->nodes[next_sensor.dest].type != NODE_SENSOR) {
      MyNextLandmark(tm, prev, curr, &next);
      next_sensor.dest = next.dest;
      next_sensor.distance += next.distance;
      prev = curr;
      curr = next.dest;
    }
    next_sensors[i] = next_sensor;
    // HACK: force out of sensor range
    next_sensor.dest = 40;
  }
}

void MyUpdateSwitchState(
    track_model *tm, int switch_number, int switch_direction) {
  int switch_id = SwitchToInternal(switch_number);
  tm->nodes[switch_id].sw.set = switch_direction;
}

int MyGetSwitchState(track_model *tm, int switch_number) {
  int switch_id = SwitchToInternal(switch_number);
  return tm->nodes[switch_id].sw.set;
}

void TrackManager() {
  int train_last_speed[NUM_TRAINS+1];
  int train_speed[NUM_TRAINS+1];
  int train_waiters[NUM_TRAINS+1];
  int has_reversed[NUM_TRAINS+1];
  int reverse_waiters[NUM_TRAINS+1];
  int i;
  for (i = 0; i <= NUM_TRAINS; i++) {
    train_last_speed[i] = 0;
    train_speed[i] = 0;
    train_waiters[i] = UNDEFINED;
    has_reversed[i] = FALSE;
    reverse_waiters[i] = UNDEFINED;
  }
  
  // state of track
  track_model tm;
  
  int sender_tid;
  
  // message and reply
  track_req msg;
  track_reply reply;

  int parent_tid;
  char track;
  Receive(&parent_tid, &track, 1);
  Reply(parent_tid, NULL, 0);

  TrackModelInit(&tm, track);
  
  RegisterAs("track-manager");
  
  for ( ; ; ) {
    Receive(&sender_tid, (char*)&msg, sizeof(track_req));
    switch (msg.type) {
    case MSG_TRACK_NEXT_LANDMARK:
      reply.result = MyNextLandmark(&tm, msg.arg0, msg.arg1, &reply.next);
      Reply(sender_tid, (char*)&reply, sizeof(track_reply));
      break;
    case MSG_TRACK_NEXT_SENSORS:
      MyNextSensors(&tm, msg.arg0, msg.arg1, msg.arg2, msg.next_sensors);
      Reply(sender_tid, NULL, 0);
      break;
    case MSG_TRACK_GET_SPEED:
      reply.result = train_speed[msg.arg0];
      Reply(sender_tid, (char*)&reply, sizeof(track_reply));
      break;
    case MSG_TRACK_UPDATE_SPEED:
      train_speed[msg.arg0] = msg.arg1;
      if ((train_speed[msg.arg0] != train_last_speed[msg.arg0]) &&
          (train_waiters[msg.arg0] != UNDEFINED)) {
        train_last_speed[msg.arg0] = train_speed[msg.arg0];
        reply.result = train_speed[msg.arg0];
        int reply_tid = train_waiters[msg.arg0];
        train_waiters[msg.arg0] = UNDEFINED;
        Reply(reply_tid, (char*)&reply, sizeof(track_reply));
      }
      Reply(sender_tid, NULL, 0);
      break;
    case MSG_TRACK_WAIT_SPEED:
      if (train_speed[msg.arg0] != train_last_speed[msg.arg0]) {
        train_last_speed[msg.arg0] = train_speed[msg.arg0];
        reply.result = train_speed[msg.arg0];
        Reply(sender_tid, (char*)&reply, sizeof(track_reply));
      } else {
        train_waiters[msg.arg0] = sender_tid;
      }
      break;
    case MSG_TRACK_SENSOR_NEXT_LANDMARK:
      MySensorNextLandmark(&tm, msg.arg0, msg.arg1, &reply.next);
      Reply(sender_tid, (char*)&reply, sizeof(track_reply));
      break;
    case MSG_TRACK_UPDATE_SWITCH:
      MyUpdateSwitchState(&tm, msg.arg0, msg.arg1);
      Reply(sender_tid, NULL, 0);
      break;
    case MSG_TRACK_GET_SWITCH:
      reply.result = MyGetSwitchState(&tm, msg.arg0);
      Reply(sender_tid, (char*)&reply, sizeof(track_reply));
      break;
    case MSG_TRACK_SWITCH_STRAIGHT:
      if (tm.nodes[msg.arg0].type != NODE_SWITCH) {
        reply.result = FALSE;
      }
      else {
        reply.result = tm.nodes[msg.arg0].sw.ahead[1].dest == msg.arg1;
      }
      Reply(sender_tid, (char*)&reply, sizeof(track_reply));
      break;
    case MSG_TRACK_SWITCH_CURVED:
      if (tm.nodes[msg.arg0].type != NODE_SWITCH) {
        reply.result = FALSE;
      }
      else {
        reply.result = tm.nodes[msg.arg0].sw.ahead[0].dest == msg.arg1;
      }
      Reply(sender_tid, (char*)&reply, sizeof(track_reply));
      break;
    case MSG_TRACK_SWITCH_AFTER_SWITCH:
      if (tm.nodes[tm.nodes[msg.arg0].sw.behind.dest].type == NODE_SWITCH) {
        if (tm.nodes[tm.nodes[msg.arg0].sw.behind.dest].sw.ahead[0].dest == msg.arg0) {
          reply.result = tm.nodes[msg.arg0].sw.behind.dest;
          reply.dir = 'C';
        }
        else if (tm.nodes[tm.nodes[msg.arg0].sw.behind.dest].sw.ahead[1].dest == msg.arg0) {
          reply.result = tm.nodes[msg.arg0].sw.behind.dest;
          reply.dir = 'S';
        }
        else
          reply.result = FALSE;
      }
      else
        reply.result = FALSE;
        
      Reply(sender_tid, (char*)&reply, sizeof(track_reply));      
      break;
    case MSG_TRACK_SHOULD_REVERSE:
      reply.result = ((tm.nodes[msg.arg1].sw.ahead[0].dest == msg.arg0 && tm.nodes[msg.arg1].sw.ahead[1].dest == msg.arg2) ||
        (tm.nodes[msg.arg1].sw.ahead[1].dest == msg.arg0 && tm.nodes[msg.arg1].sw.ahead[0].dest == msg.arg2));
        
      Reply(sender_tid, (char*)&reply, sizeof(track_reply));     
      break;
    case MSG_TRACK_GET_NODE:
      reply.node = tm.nodes[msg.arg0];
      Reply(sender_tid, (char*)&reply, sizeof(track_reply));
      break;
    default:
      Panic("TrackManager", "invalid message type");
      break;
    }
  }
  Panic("TrackManager", "Chuck Norris!");
}

int GetNextLandmark(int prev, int curr, track_edge *next) {
  int track_manager_tid = WhoIs("track-manager");
  track_req msg;
  msg.type = MSG_TRACK_NEXT_LANDMARK;
  msg.arg0 = prev;
  msg.arg1 = curr;
  track_reply reply;
  Send(track_manager_tid, (char*)&msg, sizeof(track_req),
      (char*)&reply, sizeof(track_reply));
  *next = reply.next;
  return reply.result;
}

void GetNextSensors(int sensor_id, int sensor_direction, int n,
    track_edge *next_sensors) {
  int track_manager_tid = WhoIs("track-manager");
  track_req msg;
  msg.type = MSG_TRACK_NEXT_SENSORS;
  msg.arg0 = sensor_id;
  msg.arg1 = sensor_direction;
  msg.arg2 = n;
  msg.next_sensors = next_sensors;
  Send(track_manager_tid, (char*)&msg, sizeof(track_req), NULL, 0);
}

void UpdateTrainSpeed(int train_number, int train_speed) {
  int track_manager_tid = WhoIs("track-manager");
  track_req msg;
  msg.type = MSG_TRACK_UPDATE_SPEED;
  msg.arg0 = train_number;
  msg.arg1 = train_speed;
  Send(track_manager_tid, (char*)&msg, sizeof(track_req), NULL, 0);
}

int GetTrainSpeed(int train_number) {
  int track_manager_tid = WhoIs("track-manager");
  track_req msg;
  msg.type = MSG_TRACK_GET_SPEED;
  msg.arg0 = train_number;
  track_reply reply;
  Send(track_manager_tid, (char*)&msg, sizeof(track_req),
    (char*)&reply, sizeof(track_reply));
  return reply.result;
}

int WaitTrainSpeed(int train_number) {
  int track_manager_tid = WhoIs("track-manager");
  track_req msg;
  msg.type = MSG_TRACK_WAIT_SPEED;
  msg.arg0 = train_number;
  track_reply reply;
  Send(track_manager_tid, (char*)&msg, sizeof(track_req),
    (char*)&reply, sizeof(track_reply));
  return reply.result;
}

void GetNextDirectedLandmark(
    int sensor_id, int sensor_direction, track_edge *next) {
  int track_manager_tid = WhoIs("track-manager");
  track_req msg;
  msg.type = MSG_TRACK_SENSOR_NEXT_LANDMARK;
  msg.arg0 = sensor_id;
  msg.arg1 = sensor_direction;
  track_reply reply;
  Send(track_manager_tid, (char*)&msg, sizeof(track_req),
    (char*)&reply, sizeof(track_reply));
  *next = reply.next;
}

void UpdateSwitchState(int switch_number, int switch_direction) {
  int track_manager_tid = WhoIs("track-manager");
  track_req msg;
  msg.type = MSG_TRACK_UPDATE_SWITCH;
  msg.arg0 = switch_number;
  msg.arg1 = switch_direction;
  Send(track_manager_tid, (char*)&msg, sizeof(track_req), NULL, 0);
}

int GetSwitchState(int switch_number) {
  int track_manager_tid = WhoIs("track-manager");
  track_req msg;
  msg.type = MSG_TRACK_GET_SWITCH;
  msg.arg0 = switch_number;
  track_reply reply;
  Send(track_manager_tid, (char*)&msg, sizeof(track_req),
     (char*)&reply, sizeof(track_reply));
  return reply.result;
}

int SwitchCurved(int curr, int next) {
  int track_manager_tid = WhoIs("track-manager");
  track_req msg;
  msg.type = MSG_TRACK_SWITCH_CURVED;
  msg.arg0 = curr;
  msg.arg1 = next;
  track_reply reply;
  Send(track_manager_tid, (char*)&msg, sizeof(track_req),
     (char*)&reply, sizeof(track_reply));
  return reply.result;
}

int SwitchStraight(int curr, int next) {
  int track_manager_tid = WhoIs("track-manager");
  track_req msg;
  msg.type = MSG_TRACK_SWITCH_STRAIGHT;
  msg.arg0 = curr;
  msg.arg1 = next;
  track_reply reply;
  Send(track_manager_tid, (char*)&msg, sizeof(track_req),
     (char*)&reply, sizeof(track_reply));
  return reply.result;
}

int SwitchAfterSwitch(int sw, char *dir) {
  int track_manager_tid = WhoIs("track-manager");
  track_req msg;
  msg.type = MSG_TRACK_SWITCH_AFTER_SWITCH;
  msg.arg0 = sw;
  track_reply reply;
  Send(track_manager_tid, (char*)&msg, sizeof(track_req),
     (char*)&reply, sizeof(track_reply));
  *dir = reply.dir;
  return reply.result;
}

int ShouldReverse(int next, int oneAhead, int twoAhead) {
  int track_manager_tid = WhoIs("track-manager");
  track_req msg;
  track_reply reply;
  msg.type = MSG_TRACK_SHOULD_REVERSE;
  msg.arg0 = next;
  msg.arg1 = oneAhead;
  msg.arg2 = twoAhead;
  Send(track_manager_tid, (char*)&msg, sizeof(track_req), (char*)&reply, sizeof(track_reply));
  return reply.result;
}

void GetNode(int id, track_node *node) {
  int track_manager_tid = WhoIs("track-manager"); 
  track_req msg;
  track_reply reply;
  msg.type = MSG_TRACK_GET_NODE;
  msg.arg0 = id;
  Send(track_manager_tid, (char*)&msg, sizeof(track_req),
      (char*)&reply, sizeof(reply));
  *node = reply.node;
}
